package main

import (
	"database/sql"
	"flag"
	"fmt"
	"github.com/gin-gonic/gin"
	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"strings"
)

// Config Parsing the structure of yaml
type Config struct {
	Mysql struct {
		User     string `yaml:"user"`     // username
		Password string `yaml:"password"` // password
		Host     string `yaml:"host"`     // host
		Port     int    `yaml:"port"`     // port
		Database string `yaml:"database"` // database
	} `yaml:"mysql"` // mysql
	Port int `yaml:"port"` // port
}

// funcuion used to parse the yaml file
func parseYaml(file string) (Config, error) {
	var config Config
	yamlFile, err := ioutil.ReadFile(file)
	if err != nil {
		return config, err
	}
	err = yaml.Unmarshal(yamlFile, &config)
	if err != nil {
		return config, err
	}
	return config, err
}

func main() {

	// -f flag to parse the yaml file
	configPath := flag.String("f", "config.yaml", "config file path")

	flag.Parse()

	// parse the yaml file
	config, err := parseYaml(*configPath)
	if err != nil {
		log.Panicf("Error parsing yaml file: %s", err)
	}

	// Connect to mysql
	sprintf := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8", config.Mysql.User, config.Mysql.Password, config.Mysql.Host, config.Mysql.Port, config.Mysql.Database)
	open, err := sqlx.Open("mysql", sprintf)
	if err != nil {
		log.Panicf("Error connecting to mysql: %s", err)
	}

	// New default gin object
	r := gin.Default()
	// root route
	r.GET("/", func(c *gin.Context) {
		html := "<!DOCTYPE html>\n" +
			"<html lang=\"en\">\n" +
			"<head>\n" +
			"    <meta charset=\"UTF-8\">\n" +
			"    <title>SQL注入测试</title>\n" +
			"</head>\n" +
			"<body>\n" +
			"<h3>安全的登录<h3>" +
			"<!--Send post request to gin-->\n" +
			"<form action=\"/diy/sqlInjection/security/login\" method=\"post\">\n" +
			"    <input type=\"text\" name=\"username\" placeholder=\"请输入用户名\">\n" +
			"    <input type=\"text\" name=\"password\" placeholder=\"请输入密码\">\n" +
			"    <input type=\"submit\" value=\"提交\">\n" +
			"</form>\n" +
			"<h5 style=\"color: #eeeeee\">本环境每晚12:00自动重置</h5>\n" +
			"<a href=\"/diy/sqlInjection/security\">尝试安全的登录</a><br><br>\n" +
			"<a href=\"/diy/sqlInjection\">尝试不安全的登录</a><br><br>\n" +
			"<a href=\"https://gitee.com/deng_wenyi/sql-injection\">两者的源码比较</a><br><br>\n" +
			"<span>看完源码比较，你会发现普通的代码审计并没有那么复杂</span>\n" +
			"</body>\n" +
			"</html>"

		// Write HTML to response
		_, err := c.Writer.Write([]byte(html))
		if err != nil {
			return
		}
	})

	// Receive from post request
	r.POST("/login", func(c *gin.Context) {

		// Get username and password from post request
		UserName := c.PostForm("username")
		Password := c.PostForm("password")

		// handle user input
		if len(strings.TrimSpace(UserName)) == 0 || len(strings.TrimSpace(Password)) == 0 {
			c.JSON(http.StatusOK, gin.H{
				"status":  "error",
				"message": "用户名或密码不能为空",
			})
			return
		}

		// Query the database
		// Use safe query
		rows := open.QueryRow("SELECT * FROM user WHERE username = ? AND password = ?", UserName, Password)
		if rows != nil {
			var id int
			var username string
			var password string
			var email string
			var age int
			var money float64
			err = rows.Scan(&id, &username, &password, &email, &age, &money)
			switch err {
			case nil:
			case sql.ErrNoRows:
				c.JSON(http.StatusOK, gin.H{
					"code":    http.StatusOK,
					"message": "用户名或密码错误",
				})
			default:
				log.Printf("Error querying database: %s", err)
			}
			if username == UserName && password == Password {

				// Return login success
				c.JSON(http.StatusOK, gin.H{
					"code": 200,
					"msg":  "登录成功",
					"user": gin.H{
						"id":       id,
						"username": username,
						"password": password,
						"email":    email,
						"age":      age,
						"money":    money,
					},
				})
			}
		}
	})
	err = r.Run(":" + strconv.Itoa(config.Port))
	if err != nil {
		log.Panicf("Error running gin: %s", err)
	}
}
